import numpy as np
import pandas as pd
import pickle
from model_generator.density import WFEstDensity
from model_generator.regression import WFEstReg


class Trainer:
    def __init__(self, df_name, sample_size, x_header, y_header=None, bandwidth=None):
        self.df_name = df_name
        self.sample_size = sample_size
        self.x_header = x_header
        self.y_header = y_header

        self.bandwidth = bandwidth

    def train(self):
        print('start training....')
        data = pd.read_table('../data/{}'.format(self.df_name), sep='|')
        total_size = len(data)
        sample_frac = float(self.sample_size) / float(total_size)
        print('total size: {}, sample_size: {}, sample_frac: {}'.format(total_size, self.sample_size, sample_frac))

        # 训练kde
        print('train kde for column: {}'.format(self.x_header))
        sample = data.sample(frac=sample_frac)
        kde_x = np.array(sample[self.x_header])
        kde_x = kde_x[:, np.newaxis]
        kde = WFEstDensity(bandwidth=self.bandwidth).fit(kde_x)

        # 训练回归
        data.sort_values(self.x_header, inplace=True)
        data.reset_index(drop=True, inplace=True)
        sample = data.sample(frac=sample_frac)

        if self.y_header is None:
            print('train reg with x: index, y: {}'.format(self.x_header))
            reg_x = np.array(sample.index.tolist()).reshape(-1, 1)
            reg_y = np.array(sample[self.x_header])
        else:
            print('train reg with x: {}, y: {}'.format(self.x_header, self.y_header))
            reg_x = np.array(sample[self.x_header]).reshape(-1, 1)
            reg_y = np.array(sample[self.y_header])
        reg = WFEstReg().fit(reg_x, reg_y)

        prefix = self.df_name.split('.')[0]
        output = open('../models/kde_{}_{}.pkl'.format(prefix, self.x_header), 'wb')
        pickle.dump(kde, output)
        output.close()

        if self.y_header is None:
            reg_file_name = '../models/reg_{}_{}_{}.pkl'.format(prefix, 'index', self.x_header)
        else:
            reg_file_name = '../models/reg_{}_{}_{}.pkl'.format(prefix, self.x_header, self.y_header)
        output = open(reg_file_name, 'wb')
        pickle.dump(reg, output)
        output.close()


class Loader:
    def __init__(self, df_name, x_header, y_header=None):
        self.df_name = df_name
        self.x_header = x_header
        self.y_header = y_header

    def load_kde(self):
        prefix = self.df_name.split('.')[0]
        with open('../models/kde_{}_{}.pkl'.format(prefix, self.x_header), 'rb') as f:
            kde = pickle.load(f)
        return kde

    def load_reg(self):
        prefix = self.df_name.split('.')[0]
        if self.y_header is None:
            reg_file_name = '../models/reg_{}_{}_{}.pkl'.format(prefix, 'index', self.x_header)
        else:
            reg_file_name = '../models/reg_{}_{}_{}.pkl'.format(prefix, self.x_header, self.y_header)
        with open(reg_file_name, 'rb') as f:
            reg = pickle.load(f)
        return reg


if __name__ == "__main__":
    trainer = Trainer('orders1g.tbl', 150000, 'o_totalprice', bandwidth=10)
    trainer.train()
